import torch
import os
# from tqdm import tqdm

from torch.utils.data import DataLoader
from .ModelUtil import *
from .DataProcess import *

# 模型训练函数
def train(dataloader, model, loss_fn, optimizer, device):
    size = len(dataloader.dataset)
    model.train()
    correct = 0

    for batch, (X, y) in enumerate(dataloader):
        X = X.to(torch.float32)
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y.to(torch.int64))
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
    correct /= size
    print(f"Train Error: \n Accuracy: {(100*correct):>0.1f}%" )


def train_onehot(dataloader, model, loss_fn, optimizer, device):
    size = len(dataloader.dataset)
    model.train()
    correct = 0

    for batch, (X, y) in enumerate(dataloader):
        X = X.to(torch.float32)
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)
        correct += (pred.argmax(1) == y.argmax(1)).type(torch.float).sum().item()

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
    correct /= size
    print(f"Train Error: \n Accuracy: {(100*correct):>0.1f}%" )



def train_softlabel(dataloader, model, loss_fn, optimizer, device):
    size = len(dataloader.dataset)
    model.train()
    correct = 0

    for batch, (X, y) in enumerate(dataloader):
        X = X.to(torch.float32)
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred_logits = model(X)
        pred = torch.nn.functional.softmax(pred_logits, dim=1)
        loss = loss_fn(torch.log(pred), y)


        correct += (pred.argmax(1) == y.argmax(1)).type(torch.float).sum().item()

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
    correct /= size
    print(f"Train Error: \n Accuracy: {(100*correct):>0.1f}%" )






# def train(model, device, train_loader, optimizer, epoch):
#     # 相比常规的模型训练需要添加参数：args，privacy_engine
#     model.train()
#     criterion = nn.CrossEntropyLoss()
#     losses = []
#     for _batch_idx, (data, target) in enumerate(tqdm(train_loader)):
#         data, target = data.to(device), target.to(device)
#         optimizer.zero_grad()
#         output = model(data)
#         loss = criterion(output, target)
#         loss.backward()
#         optimizer.step()
#         losses.append(loss.item())

#     print(f"Train Epoch: {epoch} \t Loss: {np.mean(losses):.6f}")
















# 应用opacus的dp模型训练函数
def train_DP(dataloader, model, loss_fn, optimizer, device, privacy_engine, delta, epoch):
    size = len(dataloader.dataset)
    model.train()
    correct = 0

    for batch, (X, y) in enumerate(dataloader):
        X = X.to(torch.float32)
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y.to(torch.int64))
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
    correct /= size
    print(f"Train Error: \n Accuracy: {(100*correct):>0.1f}%" )
    epsilon = privacy_engine.accountant.get_epsilon(delta=delta)
    print(
        f"Train Epoch: {epoch} \t"
        # f"Loss: {np.mean(losses):.6f} "
        f"(ε = {epsilon:.2f}, δ = {delta})"
    )
    return epsilon

def evaluate(dataloader, model, loss_fn, device):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X = X.to(torch.float32)
            X, y = X.to(device), y.to(device)
            # pred = model(X)

            try:
                pred,_,_ = model(X)
            except:
                try:
                    pred,_ = model(X)
                except:
                    pred = model(X)

            # print(pred.shape,y.shape)
            test_loss += loss_fn(pred, y.to(torch.int64)).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
    return correct


# 批量训练模型
def train_models(weight_dir, data_name, model, LEARNING_RATE, epochs, transform, device, batch_size=64, num_shadowsets=50, prop_keep=0.14):
    if not os.path.isdir(weight_dir):
        os.makedirs(weight_dir)
    for i in range(80,100):
        (x_train, y_train), (x_test, y_test), train_keep_exp, test_keep_exp = globals()['load_{}'.format(data_name)](i, num_shadowsets, prop_keep=prop_keep, seed=0)
        training_data = CustomDataset(x_train, y_train, transform)
        test_data = CustomDataset(x_test, y_test, transform)
        train_dataloader = DataLoader(training_data, batch_size=batch_size)
        test_dataloader = DataLoader(test_data, batch_size=batch_size)
        if model in ['NN', 'NN_4layer']:
            ReferenceModel = globals()['create_{}_model'.format(model)](x_train.shape[1], y_train.max()+1)
        elif model == 'CNN':
            ReferenceModel = globals()['create_{}_model'.format(model)](y_train.max()+1, data_name)
        else:
            ReferenceModel = globals()['create_{}_model'.format(model)](y_train.max()+1)
        # print(ReferenceModel)
        ReferenceModel.to(device)
        loss_fn = nn.CrossEntropyLoss()
        # optimizer = torch.optim.Adam(ReferenceModel.parameters(), lr=LEARNING_RATE)
        optimizer = torch.optim.SGD(ReferenceModel.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=5e-4)
        # optimizer = torch.optim.SGD(ReferenceModel.parameters(), lr=LEARNING_RATE, momentum=0.9)
        # train_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 50, 70, 90], gamma=0.2)
        train_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 75], gamma=0.2)
        for t in range(epochs):
            print(f"Epoch {t+1}\n-------------------------------")
            if t > 0:
                train_scheduler.step(t+1)
            train(train_dataloader, ReferenceModel, loss_fn, optimizer, device)
        print("Done!") 
        # evaluate(test_dataloader, ReferenceModel, loss_fn, device)
        weight_path = os.path.join(weight_dir, "{}_{}_epoch{}_shadownum{}_model{}.pth".format(data_name, model, epochs, num_shadowsets, i))
        #weight_path = os.path.join(weight_dir, "{}_{}_epoch{}_model{}.pth".format(data_name, model, epochs, i))
        torch.save(ReferenceModel.state_dict(), weight_path)



def noise_infer(dataloader, model, loss_fn, device, sigma=0.02):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    i = 0
    with torch.no_grad():
        for X, y in dataloader:
            X = X.to(torch.float32)
            X = X + sigma * torch.randn(X.shape)
            # X = torch.clip(a,-1, 1)
            X, y = X.to(device), y.to(device)
            # pred = model(X)

            try:
                pred,_,_ = model(X)
            except:
                try:
                    pred,_ = model(X)
                except:
                    pred = model(X)


            test_loss += loss_fn(pred, y.to(torch.int64)).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            result = (pred.argmax(1) == y).type(torch.float)
            if i==0:
                pred_result = result
                i = i+1
            else:
                pred_result = torch.cat((pred_result, result), 0)
                i = i+1

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
    return pred_result
